# -*- encoding: utf-8 -*-
"""
@File    : keras_layer_in_layer.py
@Author  : lilong
@Time    : 2023/2/24 10:44 上午
"""

import tensorflow as tf

from tensorflow.keras.layers import *
from tensorflow.keras import backend as K


def test_lambda():
    """对输入的矩阵进行简单操作"""
    x_in = Input(shape=(10,))
    x = Lambda(lambda x: x + 2)(x_in)  # 对输入加上2


def noise_test():
    """加噪声"""

    def add_noise_in_train(x):
        """区分训练阶段和测试阶段，比如训练阶段给输入加入一些噪声，
        而测试阶段则去掉噪声，这需要用K.in_train_phase实现，
        Lambda层仅仅适用于不需要增加训练参数的情形，如果想要实现的功能需要往模型新增参数，
        那么就必须要用到自定义Layer了
        """

        # 训练阶段，加上标准高斯噪声
        x_ = x + K.random_normal(shape=K.shape(x))
        return K.in_train_phase(x_, x)

    x_in = Input(shape=(10,))
    x = Lambda(add_noise_in_train)(x_in)  # 训练阶段加入高斯噪声，测试阶段去掉


class MyLayer(Layer):
    """自定义层
    添加可训练的参数
    """

    def __init__(self, output_dim, **kwargs):
        super(MyLayer, self).__init__(**kwargs)  # 必须
        self.output_dim = output_dim  # 可以自定义一些属性，方便调用
        self.kernel = None  # 可训练参数

    def build(self, input_shape):
        """添加可训练参数"""
        self.kernel = self.add_weight(name='kernel',
                                      shape=(input_shape[1], self.output_dim),
                                      initializer='uniform',
                                      trainable=True)

    def call(self, x, **kwargs):
        """定义功能，相当于Lambda层的功能函数"""
        return K.dot(x, self.kernel)

    def compute_output_shape(self, input_shape):
        """计算输出形状，如果输入和输出形状一致，那么可以省略，否则最好加上"""
        return input_shape[0], self.output_dim


class SplitVector(Layer):
    """双输出层

    example:
        x1, x2 = SplitVector()(x_in)
    """

    def __init__(self, **kwargs):
        super(SplitVector, self).__init__(**kwargs)

    def compute_output_shape(self, input_shape):
        """output_shape也要是对应的list"""
        in_dim = input_shape[-1]
        return (None, in_dim // 2), (None, in_dim - in_dim // 2)

    def call(self, inputs, **kwargs):
        """按第二个维度对tensor进行切片，返回一个list"""
        in_dim = K.int_shape(inputs)[-1]
        return [inputs[:, :in_dim // 2], inputs[:, in_dim // 2:]]


class OurLayer(Layer):
    """定义新的Layer，增加reuse方法，允许在定义Layer时调用现成的层
    """

    def reuse(self, layer, *args, **kwargs):
        if not layer.built:
            if len(args) > 0:
                inputs = args[0]
            else:
                inputs = kwargs['inputs']
            if isinstance(inputs, list):
                input_shape = [K.int_shape(x) for x in inputs]
            else:
                input_shape = K.int_shape(inputs)
            layer.build(input_shape)
        outputs = layer.call(*args, **kwargs)

        for w in layer.trainable_weights:
            if w not in self._trainable_weights:
                self._trainable_weights.append(w)
        for w in layer.non_trainable_weights:
            if w not in self._non_trainable_weights:
                self._non_trainable_weights.append(w)
        for u in layer.updates:
            if not hasattr(self, '_updates'):
                self._updates = []
            if u not in self._updates:
                self._updates.append(u)

        return outputs


class OurDense(OurLayer):
    """原来是继承Layer类，现在继承OurLayer类
    """

    def __init__(self, hidden_dim, output_dim,
                 hidden_activation='linear',
                 output_activation='linear',
                 **kwargs):
        super(OurDense, self).__init__(**kwargs)
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.hidden_activation = hidden_activation
        self.output_activation = output_activation

    def build(self, input_shape):
        """在build方法里边添加需要重用的层，
        当然也可以像标准写法一样条件可训练的权重。
        """
        super(OurDense, self).build(input_shape)
        self.h_dense = Dense(self.hidden_dim, activation=self.hidden_activation)
        self.o_dense = Dense(self.output_dim, activation=self.output_activation)

    def compute_output_shape(self, input_shape):
        return input_shape[:-1] + (self.output_dim,)

    def call(self, inputs, **kwargs):
        """直接reuse一下层，等价于o_dense(h_dense(inputs))"""
        h = self.reuse(self.h_dense, inputs)
        o = self.reuse(self.o_dense, h)
        return o

